package com.fwmagic.spark.ml.similarity

import com.fwmagic.spark.util.SparkUtils
import org.apache.spark.sql.types.{DataTypes, StructType}
import org.apache.spark.sql.{DataFrame, SparkSession}

import scala.collection.mutable

/**
 * id,f1,f2,f3,f4,f5
 * 1,40,50,10,2,5
 * 2,80,100,20,4,10
 * 3,121,148,28,6,15
 * 4,35,45,8,2,4
 * 5,70,92,16,4,7.9
 * 6,103,136,23,6.6,12.5
 *
 * 计算上面所有向量两两之间的余弦相似度和欧式距离衡量的相似度
 *
 */
object SimilarityDemo {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkUtils.getSparkSession(this.getClass.getSimpleName)

    val schema: StructType = new StructType()
      .add("id", DataTypes.DoubleType)
      .add("f1", DataTypes.DoubleType)
      .add("f2", DataTypes.DoubleType)
      .add("f3", DataTypes.DoubleType)
      .add("f4", DataTypes.DoubleType)
      .add("f5", DataTypes.DoubleType)

    //加载数据
    val sampleDF: DataFrame = spark.read
      .schema(schema)
      .option("header", true)
      .csv("data/similarity/demo.csv")

    //将数据注册成一张表
    sampleDF.createOrReplaceTempView("sample")

    //将函数注册到sparksql中
    //注册欧式距离相似度的函数
    spark.udf.register("eud", eud)
    //注册余弦相似度函数
    spark.udf.register("cos", cos)

    //sql方式计算
    val df: DataFrame = spark.sql(
      """
        |select
        |a.id,
        |b.id,
        |eud(array(a.f1,a.f2,a.f3,a.f4,a.f5),array(b.f1,b.f2,b.f3,b.f4,b.f5)) edu,
        |cos(array(a.f1,a.f2,a.f3,a.f4,a.f5),array(b.f1,b.f2,b.f3,b.f4,b.f5)) cos
        |from sample a join sample b on a.id < b.id
        |""".stripMargin)

    df.show(100, false)

    spark.close()
  }

  //计算欧式距离相似度的函数（距离越近，越相似）
  val eud = (f1: mutable.WrappedArray[Double], f2: mutable.WrappedArray[Double]) => {
    // [(10,20),(20,30)]
    val d2: Double = f1.zip(f2).map(t => Math.pow(t._1 - t._2, 2)).sum
    //+1是为了防止分母为0，欧式距离的导数即计算欧式距离相似度
    1 / (Math.pow(d2, 0.5) + 1)
  }

  //计算余弦相似度（余弦角度越接近，相似度越接近）
  val cos = (f1: mutable.WrappedArray[Double], f2: mutable.WrappedArray[Double]) => {
    val f1M: Double = Math.pow(f1.map(Math.pow(_, 2)).sum, 0.5)
    val f2M: Double = Math.pow(f2.map(Math.pow(_, 2)).sum, 0.5)

    val dj: Double = f1.zip(f2).map(t => (t._1 * t._2)).sum

    dj / (f1M * f2M)
  }
}
